# 结构化网格程序

import numpy as np
from precision_data import fp

class StructureMesh:
    def __init__(self):
        pass

    def CreateMesh(self, dim, ncx, ncy, ncz=1):
        # 空间维度
        self.dim = dim

        # 在每个方向上的单元数量
        self.ncx = ncx
        self.ncy = ncy
        self.ncz = ncz

        if self.dim == 2:
            self.ncz = 1

        # 每个方向上的节点数量
        self.nx = self.ncx + 1
        self.ny = self.ncy + 1
        self.nz = self.ncz + 1

        print('nx, ny, nz = ', self.nx, self.ny, self.nz)
        print('ncx, ncy, ncz = ', self.ncx, self.ncy, self.ncz)

    def CreateCoordinates(self, xmin, xmax, ymin, ymax, zmin=fp(0.0), zmax=fp(1.0)):

        # 域坐标系
        self.xmin, self.xmax = xmin, xmax
        self.ymin, self.ymax = ymin, ymax
        self.zmin, self.zmax = zmin, zmax

        if self.dim==2:
            self.zmin = fp(0.0)
            self.zmax = fp(1.0)

        # 网格节点坐标
        self.x = np.zeros(self.nx, dtype=fp)
        self.y = np.zeros(self.ny, dtype=fp)
        self.z = np.zeros(self.nz, dtype=fp)

        # 单元中心网格坐标
        self.xc = np.zeros(self.ncx, dtype=fp)
        self.yc = np.zeros(self.ncy, dtype=fp)
        self.zc = np.zeros(self.ncz, dtype=fp)

        # 网格生成
        self.dx = (self.xmax - self.xmin)/float(self.nx - 1)
        for i in range(self.nx):
            self.x[i] = self.xmin + self.dx * float(i)
        self.dy = (self.ymax - self.ymin)/float(self.ny - 1)
        for i in range(self.ny):
            self.y[i] = self.ymin + self.dy * float(i)
        self.dz = (self.zmax - self.zmin)/float(self.nz - 1)
        for i in range(self.nz):
            self.z[i] = self.zmin + self.dz * float(i)

        for i in range(self.ncx):
            self.xc[i] = 0.5 * (self.x[i] + self.x[i+1])
        for i in range(self.ncy):
            self.yc[i] = 0.5 * (self.y[i] + self.y[i+1])
        for i in range(self.ncz):
            self.zc[i] = 0.5 * (self.z[i] + self.z[i+1])
        
        print('dx, dy, dz = ', self.dx, self.dy, self.dz)
        print('bbox xmin, xmax = ', xmin, xmax)
        print('bbox ymin, ymax = ', ymin, ymax)
        print('bbox zmin, zmax = ', zmin, zmax)
        print('bbox Lx, Ly, Lz = ', xmax - xmin, ymax - ymin, zmax - zmin)

    def CreateFieldMeshData(self):

        # 定义场量--这里为温度
        self.t = np.zeros((self.ncx, self.ncy, self.ncz), dtype=fp)

        # Old temperature field
        self.t0 = np.zeros((self.ncx, self.ncy, self.ncz), dtype=fp)

    def SetInitialT(self, T):
        self.t = T * np.ones((self.ncx, self.ncy, self.ncz), dtype=fp)

    def CreateCoeffMeshData(self):

        # 设置系数的数量
        if self.dim==2:
            self.ncoef = 6  # aP, aE, aW, aN, aS, bsrc
        elif self.dim==3:
            self.ncoef = 8  # aP, aE, aW, aN, aS, aT, aB, bsrc

        # 系数存储在4维数组(索引)
        self.id_aP = 0
        self.id_aE = 1
        self.id_aW = 2
        self.id_aN = 3
        self.id_aS = 4
        if self.dim==3:
            self.id_aT = 5
            self.id_aB = 6
        self.id_bsrc = self.ncoef - 1

        # 温度系数
        self.ct = np.zeros((self.ncx, self.ncy, self.ncz , self.ncoef), dtype=fp)

    def CreateSimulationData(self):

        # equations
        self.ieqn = 0
        self.eqn_conduction = 0 
        self.eqn_flow = 1
        self.eqn_conduction_flow = 2

        # Number of non-linear iterations(针对于整个求解方程的)
        self.nsteps = 1

        # Time step, do not set it to be 0.0
        self.dt = fp(1.0)

        # Simulation loop control variables
        self.stop_sim = False

    def Set_nsteps(self, nsteps):
        self.nsteps = nsteps   

    def Set_dt(self, dt):
        self.dt = dt

    def CreateSolvingMethodData(self):

        # Number of linear solver iterations for  temperature
        self.niter_t = 10

        # Relaxation factor
        self.relax_t = fp(0.75)

        # Convergence tolerances of the linear solver iterations(针对于求解的变量的)
        self.res_t = fp(1.e-2)

        # Total number of linear iterations
        self.total_linsol_iters = 0

        # Convergence tolerances of the non-linear solver iterations(针对于整个求解方程的)
        self.temp_tol = fp(1.e-6)

        # L2-norm of residual of non-linear iterations
        self.l2_curr = fp(0.0)
        self.l2_max  = fp(-1.e20)
        self.l2_max_t = fp(-1.e20)
        self.l2_t = fp(0.0)

    def Set_temp_solver_param(self, niter_t, relax_t, res_t, temp_tol):
        self.niter_t = niter_t
        self.relax_t = relax_t
        self.res_t   = res_t
        self.temp_tol = temp_tol

class Fluid:
    def __init__(self, ncx, ncy, ncz):
        
        self.ncx = ncx
        self.ncy = ncy
        self.ncz = ncz

        # density and velocity feilds
        self.dens = np.ones((ncx, ncy, ncy), dtype=fp)
        self.mu   = np.ones((ncx, ncy, ncy), dtype=fp)

        # conduction coefficient and specific heat
        self.con  = fp(0.0)
        self.spht = fp(0.0)
        # self.expan = fp(0.0034722)

        # heat source
        self.heat_src = fp(0.0) 

    def SetInitialDenMu(self, dens, mu):

        # density and viscosity feild
        self.dens = dens * np.ones((self.ncx, self.ncy, self.ncz), dtype=fp)
        self.mu   = mu   * np.ones((self.ncx, self.ncy, self.ncz), dtype=fp)

    def Set_con_spht(self, con, spht):
        self.con  = con
        self.spht = spht

class BoundaryCondition: # None, Wall, Inlet, Outlet
    def __init__(self):
        self.type = 0

class BoundaryConditionTemp:
    def __init__(self):
        self.temp_type = 0
        self.t         = fp(0.0)
        self.heat_flux = fp(0.0)

class FluidBoundary:
    def __init__(self,dim):
        
        # face index
        self.fid_e = 0
        self.fid_w = 1
        self.fid_n = 2
        self.fid_s = 3
        self.fid_t = 4
        self.fid_b = 5

        # Physical boundary conditions types
        self.bc_none   = 0
        self.bc_wall   = 1
        self.bc_inlet  = 2
        self.bc_outlet = 3

        self.bcid_none = 0    # 内部边/面
        self.bcid_xmin = 1
        self.bcid_xmax = 2
        self.bcid_ymin = 3
        self.bcid_ymax = 4
        self.bcid_zmin = 5
        self.bcid_zmax = 6
        self.num_bcs   = 7

        self.bcs = []
        self.bcs_temp = []
        for _ in range(self.num_bcs):
            self.bcs.append(BoundaryCondition())
            self.bcs_temp.append(BoundaryConditionTemp())

        # Values of physical temperature boundary conditions
        self.temp_east    = fp(0.0)
        self.temp_west    = fp(0.0)
        self.temp_north   = fp(0.0)
        self.temp_south   = fp(0.0)
        self.temp_top     = fp(0.0)
        self.temp_bottom  = fp(0.0)

        # temperature boundary condition types
        self.temp_bc_constant = 0
        self.temp_bc_heatflux = 1



        

class PostProcessing:
    def __init__(self):
        
        # Output file names
        self.vtk_fname_temp = "post_temp.vtk"

    def WriteVTKCollocated_temp(self, case):

        # Build the local data for np array
        nx = case.nx
        ny = case.ny
        nz = case.nz

        x = case.x
        y = case.y
        z = case.z

        t = case.t  

        # load "conduction_plate.plot"
        # # Open temperature output files
        with open('tem_yline.dat','w') as file3:

            # Write temperature data along y-line
            i1 = int(fp(0.0833/0.833)*nx)
            i2 = int(fp(0.5/0.833)*nx)
            k = 0
            for j in range(ny-1):
                file3.write(f"{(y[j] + y[j+1]) * 0.5} {t[i1, j , k]} {t[i2, j, k]}\n")

        # Open VTK output file
        with open(self.vtk_fname_temp, 'w') as vtk_fid:

            # Write header 
            vtk_fid.write('# vtk DataFile Version 3.0\n')
            vtk_fid.write('flash 3d grid and solution\n')
            vtk_fid.write('ASCII\n')
            vtk_fid.write('DATASET RECTILINEAR_GRID\n')

            # Write mesh information(节点)
            vtk_fid.write(f"DIMENSIONS {nx} {ny} {nz}\n")
            vtk_fid.write(f"X_COORDINATES {nx} float\n")
            vtk_fid.write(' '.join(str(i) for i in x) + '\n')
            vtk_fid.write(f"Y_COORDINATES {ny} float\n")
            vtk_fid.write(' '.join(str(i) for i in y) + '\n')
            vtk_fid.write(f"Z_COORDINATES {ny} float\n")
            vtk_fid.write(' '.join(str(i) for i in z) + '\n')

            # Write cell data
            ncell = (nx - 1) * (ny - 1 ) * (nz - 1)
            vtk_fid.write(f"CELL_DATA {ncell}\n")

            vtk_fid.write('{:s}'.format("FIELD FieldData 1\n"))

            # Write temperature data
            vtk_fid.write(f"t 1 {ncell} float\n")
            t_arr = np.ravel(t[:, :, :], order='F')
            vtk_fid.write(' '.join(str(i) for i in t_arr) + '\n')



        